fix(scheduler): read the prefill chunk cap from the pool, not a snapshot - #439
gberasmus87 wants to merge 1 commit into
Conversation
`CacheManager.__init__` copied `swa_pool.prefill_chunk_budget` into an instance attribute, so the cap kept its construction-time value for the life of the manager and `rebuild` never refreshed it. `Scheduler.rebuild_cache` recomputes `prefill_budget` from `cache_manager.prefill_chunk_budget` after a runtime resize, so reading a frozen value meant it wrote back the number it already had. Two consequences: * Growing the pool leaves prefill chunking where it was. Measured on DSV4-Flash through `POST /v1/cache/rebuild`: the window pool went 100 -> 215 pages and the chunk budget stayed at 4864, so an 11.7k prompt still took three whole-layer expert streams (32.5 s) instead of the one it had just been sized for. * Shrinking is worse, and is the hazard `rebuild_cache`'s own comment warns about: the stale cap is then too LARGE for the pool, so the next long prompt is chunked past what `_alloc_window` can satisfy. Making it a property that re-reads the pool on every access fixes both. Nothing assigns to it -- `scheduler.py` reads it in two places and `DSV4PagedKVCache` owns the value -- so a read-only property is the whole change. `test_rebuild_cache_refreshes_prefill_budget` already covers the scheduler half, but with a `SimpleNamespace` manager whose cap the test sets by hand, so it cannot see this: the gap is between a real manager and its real pool. The new test closes exactly that gap. It drives `DSV4PagedKVCache.rebuild` (whose `_init_paged_state` recomputes `_chunk_budget` from the new window-slot count) rather than assigning the attribute, and checks both directions -- grown and shrunk -- asserting the pool's cap actually moved first so the test cannot pass vacuously. Verified: the new test fails on the unfixed tree and passes with the fix; the rest of tests/scheduler/ is unchanged either way.
|
Thanks for this — the frozen-snapshot read is exactly right, and the two-direction test (driving Wanted to align on related ground we've measured on DSV4-Flash (TP8, 8×RTX 4090, ~330k-token pool), in case it affects what this PR is expected to cover: The chunk budget is not only stale across
So we ended up pinning the effective chunk with the existing One more thing your Happy to share the branch if the geometric-decoupling direction is interesting to maintainers. |
This branch has not been deployed
2 participants
fix(scheduler): read the prefill chunk cap from the pool, not a snapshot
CacheManager.__init__copiedswa_pool.prefill_chunk_budgetinto an instanceattribute, so the cap kept its construction-time value for the life of the
manager and
rebuildnever refreshed it.Scheduler.rebuild_cacherecomputesprefill_budgetfromcache_manager.prefill_chunk_budgetafter a runtime resize, so reading a frozenvalue meant it wrote back the number it already had. Two consequences:
DSV4-Flash through
POST /v1/cache/rebuild: the window pool went 100 -> 215pages and the chunk budget stayed at 4864, so an 11.7k prompt still took three
whole-layer expert streams (32.5 s) instead of the one it had just been sized
for.
rebuild_cache's own comment warnsabout: the stale cap is then too LARGE for the pool, so the next long prompt is
chunked past what
_alloc_windowcan satisfy.Making it a property that re-reads the pool on every access fixes both. Nothing
assigns to it --
scheduler.pyreads it in two places andDSV4PagedKVCacheowns the value -- so a read-only property is the whole change.
test_rebuild_cache_refreshes_prefill_budgetalready covers the scheduler half,but with a
SimpleNamespacemanager whose cap the test sets by hand, so itcannot see this: the gap is between a real manager and its real pool. The new
test closes exactly that gap. It drives
DSV4PagedKVCache.rebuild(whose_init_paged_staterecomputes_chunk_budgetfrom the new window-slot count)rather than assigning the attribute, and checks both directions -- grown and
shrunk -- asserting the pool's cap actually moved first so the test cannot pass
vacuously.
Verified: the new test fails on the unfixed tree and passes with the fix; the
rest of tests/scheduler/ is unchanged either way.